{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d5f19d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from math import sqrt\n",
    "import copy\n",
    "import  traceback\n",
    "import shutil\n",
    "import random\n",
    "\n",
    "import numpy as np  # linear algebra\n",
    "import pydicom\n",
    "from pydicom.errors import InvalidDicomError\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "from pydicom.uid import UID\n",
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "import openpyxl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ec30c90",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_scan(path):\n",
    "    slices = [] #slices = [pydicom.dcmread(path + '/' + s) for s in filter(lambda x: x.endswith('.dcm'), os.listdir(path))]\n",
    "    for s in os.listdir(path):\n",
    "        if os.path.isdir(os.path.join(path, s)): #if not s.endswith('.dcm'):\n",
    "            continue\n",
    "        sl = pydicom.dcmread(os.path.join(path, s), force=True)\n",
    "        try:\n",
    "            sl_p = sl.pixel_array\n",
    "        except (AttributeError, InvalidDicomError):\n",
    "            traceback.print_exc()\n",
    "            print(f'\\tDelete {os.path.join(path, s)}')\n",
    "            os.remove(os.path.join(path, s))\n",
    "        else:\n",
    "            slices.append(sl)\n",
    "    slices.sort(key=lambda x: float(x.InstanceNumber))\n",
    "    return slices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaa0f64c",
   "metadata": {},
   "outputs": [],
   "source": [
    "window_width, window_level = 600, 200\n",
    "lower_b, upper_b = window_level - window_width//2, window_level + window_width//2\n",
    "print(lower_b, upper_b)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "883aa1c6",
   "metadata": {},
   "source": [
    "# 1.阴性数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2720e77f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 打印哪个病例没有2\n",
    "def print_no_cta(input_dir):\n",
    "    print(f'**********{input_dir}')\n",
    "    no_cta_list = []\n",
    "    for patient in sorted(os.listdir(input_dir)):\n",
    "        patient_path = os.path.join(input_dir, patient)\n",
    "        if os.path.isfile(patient_path): continue\n",
    "        if '2' not in os.listdir(patient_path):\n",
    "            no_cta_list.append(patient_path)\n",
    "            print(patient_path, os.listdir(patient_path))\n",
    "            continue\n",
    "        if f'images_{lower_b}_{upper_b}' not in os.listdir(os.path.join(patient_path, '2')):\n",
    "            print(f'have 2 but not have images_{lower_b}_{upper_b}', patient_path)\n",
    "    return no_cta_list\n",
    "            \n",
    "no_cta_list = []\n",
    "no_cta_list.extend(print_no_cta('/nfs3-p2/zsxm/dataset/2021-9-17-negative'))\n",
    "no_cta_list.extend(print_no_cta('/nfs3-p2/zsxm/dataset/2021-9-29-negative'))\n",
    "print(no_cta_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f310245b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#将某个scan重命名为2，如果thickness距离1的thickness相同则选择thickness小的重命名\n",
    "for patient in no_cta_list:\n",
    "    scans = os.listdir(patient)\n",
    "    if '1' not in scans:\n",
    "        print(patient, 'not have 1')\n",
    "        continue\n",
    "    if len(scans) == 2:\n",
    "        for scan in scans:\n",
    "            if scan != '1':\n",
    "                os.rename(os.path.join(patient, scan), os.path.join(patient, '2'))\n",
    "    else:\n",
    "        tk_list = []\n",
    "        for scan in scans:\n",
    "            for s in os.listdir(os.path.join(patient, scan)):\n",
    "                if os.path.isdir(os.path.join(patient, scan, s)) or not s.endswith('.dcm'):\n",
    "                    continue\n",
    "                sl = pydicom.dcmread(os.path.join(patient, scan, s))\n",
    "                try:\n",
    "                    sl_p = sl.pixel_array\n",
    "                except AttributeError:\n",
    "                    continue\n",
    "                else:\n",
    "                    if scan == '1':\n",
    "                        ct_thickness = sl.SliceThickness\n",
    "                    else:\n",
    "                        tk_list.append((sl.SliceThickness, scan))\n",
    "        min_dis, min_scan, min_tk = 10000, None, 10000\n",
    "        for tk, scan in tk_list:\n",
    "            dis = abs(tk-ct_thickness)\n",
    "            if dis < min_dis or (dis == min_dis and tk < min_tk):\n",
    "                min_dis, min_scan, min_tk = dis, scan, tk\n",
    "        print(patient, min_scan)\n",
    "        os.rename(os.path.join(patient, min_scan), os.path.join(patient, '2'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15ebe289",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_label(input_path):\n",
    "    workbook_path = os.path.join(input_path, 'label.xlsx')\n",
    "    wb = openpyxl.load_workbook(workbook_path)\n",
    "    sheet = wb['Sheet1']\n",
    "    for i, row in enumerate(sheet.iter_rows()):\n",
    "        if i == 0: continue\n",
    "        if row[3].value is not None:\n",
    "            lsct = row[3].value.split('-')\n",
    "            assert len(lsct) == 4, f'{input_path}:{patient} CT label wrong:{lsct}'\n",
    "            assert int(lsct[0]) < int(lsct[1]) < int(lsct[2]) < int(lsct[3]), f'{input_path}:{patient} CT label error:{lsct}'\n",
    "        if row[4].value is not None:\n",
    "            ls = row[4].value.split('-')\n",
    "            assert len(ls) == 4, f'{input_path}:{patient} CTA label wrong:{ls}'\n",
    "            assert int(ls[0]) < int(ls[1]) < int(ls[2]) < int(ls[3]), f'{input_path}:{patient} CTA label error:{ls}'\n",
    "            \n",
    "check_label('/nfs3-p1/zsxm/dataset/2021-9-17-negative/')\n",
    "check_label('/nfs3-p1/zsxm/dataset/2021-9-29-negative/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ad8ce0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 将2下的dcm文件根据窗宽窗位转化为png图片\n",
    "def generate_image(input_folder):\n",
    "    for patient in sorted(os.listdir(input_folder)):\n",
    "        if os.path.isfile(os.path.join(input_folder, patient)) or f'images_{lower_b}_{upper_b}' in os.listdir(os.path.join(input_folder, patient, '2')):\n",
    "            continue\n",
    "        print(f'****Processing {patient}****')\n",
    "        for scan in os.listdir(os.path.join(input_folder, patient)):\n",
    "            if scan != '2':\n",
    "                continue\n",
    "            name = patient #name = patient.split('-')[0]\n",
    "            image_path = os.path.join(input_folder, patient, scan, f'images_{lower_b}_{upper_b}')\n",
    "            if os.path.exists(image_path):\n",
    "                shutil.rmtree(image_path)\n",
    "            os.mkdir(image_path)\n",
    "\n",
    "            ct = load_scan(os.path.join(input_folder, patient, scan))\n",
    "            print_flag = False\n",
    "            for i in range(len(ct)):\n",
    "                img = ct[i].pixel_array.astype(np.int16)\n",
    "                intercept = ct[i].RescaleIntercept\n",
    "                slope = ct[i].RescaleSlope\n",
    "                if slope != 1:\n",
    "                    img = (slope * img.astype(np.float64)).astype(np.int16)\n",
    "                img += np.int16(intercept)\n",
    "                img = np.clip(img, lower_b, upper_b)\n",
    "                img = ((img-lower_b)/(upper_b-lower_b)*255).astype(np.uint8)\n",
    "                img = Image.fromarray(img)\n",
    "                if img.height != img.width:\n",
    "                    if not print_flag:\n",
    "                        print(patient, f'height({img.height}) not equal to width({img.width})\\n')\n",
    "                        print_flag = True\n",
    "                    height = width = min(img.height, img.width)\n",
    "                    if img.height != height:\n",
    "                        start = (img.height - height) / 2\n",
    "                        img = img.crop((0, start, img.width, start + height))\n",
    "                    elif img.width != width:\n",
    "                        start = (img.width - width) / 2\n",
    "                        img = img.crop((start, 0, start + height, img.height))\n",
    "                img.save(os.path.join(image_path, f'{name}_{i:04d}.png'))\n",
    "\n",
    "generate_image('/nfs3-p1/zsxm/dataset/2021-9-17-negative/')\n",
    "generate_image('/nfs3-p1/zsxm/dataset/2021-9-29-negative/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f1cd315",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 将各个病例中的png图片文件夹统一移动到一起供yolov5检测, not_move=True表示若有labels则不移动去检测\n",
    "def move_together_for_detect(input_folder, dst_path, not_move=True):   \n",
    "    if not os.path.exists(dst_path):\n",
    "        os.mkdir(dst_path)\n",
    "    root_name = input_folder.split('/')[-1] if input_folder.split('/')[-1] != '' else input_folder.split('/')[-2]\n",
    "    dst_path = os.path.join(dst_path, root_name)\n",
    "\n",
    "    for patient in sorted(os.listdir(input_folder)):\n",
    "        if os.path.isfile(os.path.join(input_folder, patient)):\n",
    "            continue\n",
    "        if not_move and os.path.exists(os.path.join(input_folder, patient, '2', 'labels')) \\\n",
    "        and os.path.exists(os.path.join(input_folder, patient, '2', f'pred_images_{lower_b}_{upper_b}')):\n",
    "            continue\n",
    "        print(f'****Processing {patient}****')\n",
    "        name = patient #name = patient.split('-')[0]\n",
    "        if os.path.exists(os.path.join(dst_path, name)):\n",
    "            print(f\"\\tremove {os.path.join(dst_path, name)}\")\n",
    "            shutil.rmtree(os.path.join(dst_path, name))\n",
    "        try:\n",
    "            shutil.copytree(os.path.join(input_folder, patient, '2', f'images_{lower_b}_{upper_b}'), os.path.join(dst_path, name))\n",
    "        except:\n",
    "            traceback.print_exc()\n",
    "\n",
    "move_together_for_detect('/nfs3-p1/zsxm/dataset/2021-9-17-negative/', '/nfs3-p1/zsxm/dataset/9_detect/')\n",
    "move_together_for_detect('/nfs3-p1/zsxm/dataset/2021-9-29-negative/', '/nfs3-p1/zsxm/dataset/9_detect/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86b97314",
   "metadata": {},
   "outputs": [],
   "source": [
    "#将检测结果移动回原文件夹内\n",
    "def move_back(result_path, ori_path):\n",
    "    if not os.path.exists(result_path):\n",
    "        print(f'目录不存在：{result_path}')\n",
    "        return\n",
    "    for patient in sorted(os.listdir(result_path)):\n",
    "        print(f'Processing {patient}')\n",
    "        p_res_path = os.path.join(result_path, patient)\n",
    "        o_res_path = os.path.join(ori_path, patient, '2', f'pred_images_{lower_b}_{upper_b}')\n",
    "        if os.path.exists(o_res_path):\n",
    "            shutil.rmtree(o_res_path)\n",
    "        os.mkdir(o_res_path)\n",
    "        for file in os.listdir(p_res_path):\n",
    "            if os.path.isfile(os.path.join(p_res_path, file)):\n",
    "                shutil.move(os.path.join(p_res_path, file), os.path.join(o_res_path, file))\n",
    "            elif os.path.isdir(os.path.join(p_res_path, file)):\n",
    "                if os.path.exists(os.path.join(ori_path, patient, '2', file)):\n",
    "                    shutil.rmtree(os.path.join(ori_path, patient, '2', file))\n",
    "                shutil.move(os.path.join(p_res_path, file), os.path.join(ori_path, patient, '2', file))\n",
    "        os.rmdir(p_res_path)\n",
    "    os.rmdir(result_path)\n",
    "                \n",
    "move_back('/home/zsxm/pythonWorkspace/yolov5_old/runs/detect/2021-9-17-negative', '/nfs3-p1/zsxm/dataset/2021-9-17-negative/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12b1176f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 切出主动脉,这里有问题啊，branch_end前后0.3的切片切不出来，建议以后更改为和下面一样的方案\n",
    "def find_coordinate(height, width, label_file, aorta):\n",
    "    with open(label_file, 'r') as f:\n",
    "        lines = f.readlines()\n",
    "    assert len(lines) <= 2, f'label.txt应该存储不多于2个label：{label_file.split(\"/\")[-1]}'\n",
    "    if len(lines) == 1:\n",
    "        assert aorta == 'j', f'如果只有一个label那么此时应为降主动脉, 但实际为{aorta}：{label_file.split(\"/\")[-1]}'\n",
    "        corr = list(map(lambda x: float(x), lines[0].split()))\n",
    "        x, y, w, h = corr[1], corr[2], corr[3], corr[4]\n",
    "        assert 0.25 < x < 0.75 and 0.2 < y < 0.8, f'边界框中心({x}, {y})出界：{label_file.split(\"/\")[-1]}'\n",
    "    else:\n",
    "        corr1, corr2 = list(map(lambda x: float(x), lines[0].split())), list(map(lambda x: float(x), lines[1].split()))\n",
    "        assert 0.25 < corr1[1] < 0.75 and 0.15 < corr1[2] < 0.85, f'边界框1中心({corr1[1]}, {corr1[2]})出界：{label_file.split(\"/\")[-1]}'\n",
    "        assert 0.25 < corr2[1] < 0.75 and 0.15 < corr2[2] < 0.85, f'边界框2中心({corr2[1]}, {corr2[2]})出界：{label_file.split(\"/\")[-1]}'\n",
    "        if aorta == 's':\n",
    "            x, y, w, h = (corr1[1], corr1[2], corr1[3], corr1[4]) if corr1[2] < corr2[2] else (corr2[1], corr2[2], corr2[3], corr2[4])\n",
    "        elif aorta == 'j':\n",
    "            x, y, w, h = (corr1[1], corr1[2], corr1[3], corr1[4]) if corr1[2] > corr2[2] else (corr2[1], corr2[2], corr2[3], corr2[4])\n",
    "        else:\n",
    "            raise Exception(f'aorta 应该为\"s\"或\"j\"其中之一: {label_file.split(\"/\")[-1]}')\n",
    "    w, h = int(width*w), int(height*h)\n",
    "    w, h = max(w, h), max(w, h)\n",
    "    return int(width*x-w/2), int(height*y-h/2), int(width*x+w/2+1), int(height*y+h/2+1)\n",
    "\n",
    "def crop_images(input_path, error_patient_list):\n",
    "    workbook_path = os.path.join(input_path, 'label.xlsx')\n",
    "    wb = openpyxl.load_workbook(workbook_path)\n",
    "    sheet = wb['Sheet1']\n",
    "    \n",
    "    for patient in sorted(os.listdir(input_path)):\n",
    "        if os.path.isfile(os.path.join(input_path, patient)):\n",
    "            continue\n",
    "        flag = True\n",
    "        for row in sheet.iter_rows():\n",
    "            if row[0].value == patient.split('-')[0]:\n",
    "                if row[3].value is not None and row[4].value is not None:\n",
    "                    flag = False\n",
    "                    ls = row[4].value.split('-')\n",
    "                    assert len(ls) == 4, f'{patient} ls wrong'\n",
    "                    aorta_start, branch_start = int(ls[0])-1, int(ls[1])-1\n",
    "                    branch_end, aorta_end = int(ls[2])-1, int(ls[3])-1\n",
    "                    lsct = row[3].value.split('-')\n",
    "                    assert len(lsct) == 4, f'{patient} lsct wrong'\n",
    "                    ct_start, ct_end = int(lsct[0])-1, int(lsct[3])-1\n",
    "                break\n",
    "        if flag: continue\n",
    "        print(f'******Processing {patient}******')\n",
    "        image_path = os.path.join(input_path, patient, '2', f'images_{lower_b}_{upper_b}')\n",
    "        label_path = os.path.join(input_path, patient, '2', 'labels')\n",
    "        crop_path = os.path.join(input_path, patient, '2', f'crops_{lower_b}_{upper_b}')\n",
    "        if os.path.exists(crop_path):\n",
    "            shutil.rmtree(crop_path)\n",
    "        os.mkdir(crop_path)\n",
    "        \n",
    "        crop_flag = True\n",
    "        offset = branch_end - branch_start\n",
    "        start, end = branch_start + int(0.1*offset), branch_end - int(0.2*offset)\n",
    "        for i in range(start, end):\n",
    "            img = Image.open(os.path.join(image_path, f'{patient}_{i:04d}.png'))\n",
    "            img = np.array(img)\n",
    "            try:\n",
    "                x1, y1, x2, y2 = find_coordinate(*img.shape[0:2], os.path.join(label_path, f'{patient}_{i:04d}.txt'), 's')\n",
    "            except:\n",
    "                traceback.print_exc()\n",
    "                crop_flag = False\n",
    "            else:#if crop_flag:\n",
    "                crop = img[y1:y2, x1:x2]\n",
    "                crop = Image.fromarray(crop)\n",
    "                crop.save(os.path.join(crop_path, f'{patient}_s_{i:04d}.png'))\n",
    "            try:\n",
    "                x1, y1, x2, y2 = find_coordinate(*img.shape[0:2], os.path.join(label_path, f'{patient}_{i:04d}.txt'), 'j')\n",
    "            except:\n",
    "                traceback.print_exc()\n",
    "                crop_flag = False\n",
    "            else:#if crop_flag:\n",
    "                crop = img[y1:y2, x1:x2]\n",
    "                crop = Image.fromarray(crop)\n",
    "                crop.save(os.path.join(crop_path, f'{patient}_j_{i:04d}.png'))\n",
    "        offset = aorta_end - branch_end\n",
    "        start, end = branch_end + int(0.1*offset), aorta_end - int(0.2*offset)\n",
    "        for i in range(start, end):\n",
    "            img = Image.open(os.path.join(image_path, f'{patient}_{i:04d}.png'))\n",
    "            img = np.array(img)\n",
    "            try:\n",
    "                x1, y1, x2, y2 = find_coordinate(*img.shape[0:2], os.path.join(label_path, f'{patient}_{i:04d}.txt'), 'j')\n",
    "            except:\n",
    "                traceback.print_exc()\n",
    "                crop_flag = False\n",
    "            else:#if crop_flag:\n",
    "                crop = img[y1:y2, x1:x2]\n",
    "                crop = Image.fromarray(crop)\n",
    "                crop.save(os.path.join(crop_path, f'{patient}_j_{i:04d}.png'))\n",
    "        if not crop_flag:\n",
    "            error_patient_list.append(patient)\n",
    "            \n",
    "epl1 = []\n",
    "crop_images('/nfs3-p1/zsxm/dataset/2021-9-17-negative/', epl1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6d39935",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(epl1))\n",
    "print(epl1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbd7e476",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 切出范围外冗余为3的主动脉\n",
    "def find_coordinate(height, width, label_file, aorta):\n",
    "    with open(label_file, 'r') as f:\n",
    "        lines = f.readlines()\n",
    "    assert len(lines) <= 2, f'label.txt应该存储不多于2个label：{label_file.split(\"/\")[-1]}'\n",
    "    if len(lines) == 1:\n",
    "        assert aorta == 'j', f'如果只有一个label那么此时应为降主动脉, 但实际为{aorta}：{label_file.split(\"/\")[-1]}'\n",
    "        corr = list(map(lambda x: float(x), lines[0].split()))\n",
    "        x, y, w, h = corr[1], corr[2], corr[3], corr[4]\n",
    "        assert 0.25 < x < 0.75 and 0.2 < y < 0.8, f'边界框中心({x}, {y})出界：{label_file.split(\"/\")[-1]}'\n",
    "    else:\n",
    "        corr1, corr2 = list(map(lambda x: float(x), lines[0].split())), list(map(lambda x: float(x), lines[1].split()))\n",
    "        assert 0.25 < corr1[1] < 0.75 and 0.15 < corr1[2] < 0.85, f'边界框1中心({corr1[1]}, {corr1[2]})出界：{label_file.split(\"/\")[-1]}'\n",
    "        assert 0.25 < corr2[1] < 0.75 and 0.15 < corr2[2] < 0.85, f'边界框2中心({corr2[1]}, {corr2[2]})出界：{label_file.split(\"/\")[-1]}'\n",
    "        if aorta == 's':\n",
    "            x, y, w, h = (corr1[1], corr1[2], corr1[3], corr1[4]) if corr1[2] < corr2[2] else (corr2[1], corr2[2], corr2[3], corr2[4])\n",
    "        elif aorta == 'j':\n",
    "            x, y, w, h = (corr1[1], corr1[2], corr1[3], corr1[4]) if corr1[2] > corr2[2] else (corr2[1], corr2[2], corr2[3], corr2[4])\n",
    "        else:\n",
    "            raise Exception(f'aorta 应该为\"s\"或\"j\"其中之一: {label_file.split(\"/\")[-1]}')\n",
    "    w, h = int(width*w), int(height*h)\n",
    "    w, h = max(w, h), max(w, h)\n",
    "    return int(width*x-w/2), int(height*y-h/2), int(width*x+w/2+1), int(height*y+h/2+1)\n",
    "\n",
    "def crop_images(input_path, error_patient_list):\n",
    "    workbook_path = os.path.join(input_path, 'label.xlsx')\n",
    "    wb = openpyxl.load_workbook(workbook_path)\n",
    "    sheet = wb['Sheet1']\n",
    "    \n",
    "    for patient in sorted(os.listdir(input_path)):\n",
    "        if os.path.isfile(os.path.join(input_path, patient)):\n",
    "            continue\n",
    "        flag = True\n",
    "        for row in sheet.iter_rows():\n",
    "            if row[0].value == patient.split('-')[0]:\n",
    "                if row[3].value is not None and row[4].value is not None:\n",
    "                    flag = False\n",
    "                    ls = row[4].value.split('-')\n",
    "                    assert len(ls) == 4, f'{patient} ls wrong'\n",
    "                    aorta_start, branch_start = int(ls[0])-1, int(ls[1])-1\n",
    "                    branch_end, aorta_end = int(ls[2])-1, int(ls[3])-1\n",
    "                break\n",
    "        if flag: continue\n",
    "        print(f'******Processing {patient}******')\n",
    "        image_path = os.path.join(input_path, patient, '2', f'images_{lower_b}_{upper_b}')\n",
    "        label_path = os.path.join(input_path, patient, '2', 'labels')\n",
    "        crop_path = os.path.join(input_path, patient, '2', f'crops3_{lower_b}_{upper_b}')\n",
    "        if os.path.exists(crop_path):\n",
    "            shutil.rmtree(crop_path)\n",
    "        os.mkdir(crop_path)\n",
    "        \n",
    "        crop_flag = True\n",
    "        offset = branch_end - branch_start\n",
    "        start, end = branch_start + int(0.1*offset), branch_end - int(0.2*offset)\n",
    "        for i in range(start-3, end+3):\n",
    "            try:\n",
    "                img = Image.open(os.path.join(image_path, f'{patient}_{i:04d}.png'))\n",
    "                img = np.array(img)\n",
    "                x1, y1, x2, y2 = find_coordinate(*img.shape[0:2], os.path.join(label_path, f'{patient}_{i:04d}.txt'), 's')\n",
    "            except:\n",
    "                traceback.print_exc()\n",
    "                crop_flag = False\n",
    "            else:#if crop_flag:\n",
    "                crop = img[y1:y2, x1:x2]\n",
    "                crop = Image.fromarray(crop)\n",
    "                if start <= i < end:\n",
    "                    crop.save(os.path.join(crop_path, f'{patient}_s_{i:04d}.png'))\n",
    "                else:\n",
    "                    crop.save(os.path.join(crop_path, f'{patient}_s_{i:04d}_n.png'))\n",
    "        offset = aorta_end - branch_start\n",
    "        start, end = branch_start + int(0.05*offset), aorta_end - int(0.1*offset)\n",
    "        for i in range(start-3, end+3):\n",
    "            try:\n",
    "                img = Image.open(os.path.join(image_path, f'{patient}_{i:04d}.png'))\n",
    "                img = np.array(img)\n",
    "                x1, y1, x2, y2 = find_coordinate(*img.shape[0:2], os.path.join(label_path, f'{patient}_{i:04d}.txt'), 'j')\n",
    "            except:\n",
    "                traceback.print_exc()\n",
    "                crop_flag = False\n",
    "            else:#if crop_flag:\n",
    "                crop = img[y1:y2, x1:x2]\n",
    "                crop = Image.fromarray(crop)\n",
    "                if start <= i < end:\n",
    "                    crop.save(os.path.join(crop_path, f'{patient}_j_{i:04d}.png'))\n",
    "                else:\n",
    "                    crop.save(os.path.join(crop_path, f'{patient}_j_{i:04d}_n.png'))\n",
    "        if not crop_flag:\n",
    "            error_patient_list.append(patient)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "420033c2",
   "metadata": {},
   "source": [
    "# 2.疾病数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d259d00f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 打印哪个病例没有2\n",
    "def print_no_cta(input_dir):\n",
    "    print(f'**********{input_dir}')\n",
    "    no_cta_list = []\n",
    "    for patient in sorted(os.listdir(input_dir)):\n",
    "        patient_path = os.path.join(input_dir, patient)\n",
    "        if os.path.isfile(patient_path): continue\n",
    "        if '2' not in os.listdir(patient_path):\n",
    "            no_cta_list.append(patient_path)\n",
    "            print(patient_path, os.listdir(patient_path))\n",
    "            continue\n",
    "        if f'images_{lower_b}_{upper_b}' not in os.listdir(os.path.join(patient_path, '2')):\n",
    "            print(f'have 2 but not have images_{lower_b}_{upper_b}', patient_path)\n",
    "    return no_cta_list\n",
    "\n",
    "no_cta_list = []\n",
    "no_cta_list.extend(print_no_cta('/nfs3-p1/zsxm/dataset/2021-9-8'))\n",
    "no_cta_list.extend(print_no_cta('/nfs3-p1/zsxm/dataset/2021-9-13/'))\n",
    "no_cta_list.extend(print_no_cta('/nfs3-p1/zsxm/dataset/2021-9-19/'))\n",
    "no_cta_list.extend(print_no_cta('/nfs3-p1/zsxm/dataset/2021-9-28/'))\n",
    "no_cta_list.extend(print_no_cta('/nfs3-p2/zsxm/dataset/2021-10-19-imh/'))\n",
    "no_cta_list.extend(print_no_cta('/nfs3-p2/zsxm/dataset/2021-10-19-pau/'))\n",
    "no_cta_list.extend(print_no_cta('/nfs3-p2/zsxm/dataset/2021-10-19-aa/'))\n",
    "no_cta_list.extend(print_no_cta('/nfs3-p2/zsxm/dataset/2021-11-20/'))\n",
    "no_cta_list.extend(print_no_cta('/nfs3-p2/zsxm/dataset/2021-11-20-imh/'))\n",
    "no_cta_list.extend(print_no_cta('/nfs3-p2/zsxm/dataset/2021-11-20-pau/'))\n",
    "print(no_cta_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41b1c5d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#将某个scan重命名为2，如果thickness距离1的thickness相同则选择thickness小的重命名\n",
    "for patient in no_cta_list:\n",
    "    scans = os.listdir(patient)\n",
    "    if '1' not in scans:\n",
    "        print(patient, 'not have 1')\n",
    "        continue\n",
    "    if len(scans) == 2:\n",
    "        for scan in scans:\n",
    "            if scan != '1':\n",
    "                os.rename(os.path.join(patient, scan), os.path.join(patient, '2'))\n",
    "    else:\n",
    "        tk_list = []\n",
    "        for scan in scans:\n",
    "            for s in os.listdir(os.path.join(patient, scan)):\n",
    "                if os.path.isdir(os.path.join(patient, scan, s)) or not s.endswith('.dcm'):\n",
    "                    continue\n",
    "                sl = pydicom.dcmread(os.path.join(patient, scan, s))\n",
    "                try:\n",
    "                    sl_p = sl.pixel_array\n",
    "                except AttributeError:\n",
    "                    continue\n",
    "                else:\n",
    "                    if scan == '1':\n",
    "                        ct_thickness = sl.SliceThickness\n",
    "                    else:\n",
    "                        tk_list.append((sl.SliceThickness, scan))\n",
    "        min_dis, min_scan, min_tk = 10000, None, 10000\n",
    "        for tk, scan in tk_list:\n",
    "            dis = abs(tk-ct_thickness)\n",
    "            if dis < min_dis or (dis == min_dis and tk < min_tk):\n",
    "                min_dis, min_scan, min_tk = dis, scan, tk\n",
    "        print(patient, min_scan)\n",
    "        os.rename(os.path.join(patient, min_scan), os.path.join(patient, '2'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d55e0ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_label(input_path):\n",
    "    workbook_path = os.path.join(input_path, 'label.xlsx')\n",
    "    wb = openpyxl.load_workbook(workbook_path)\n",
    "    sheet = wb['Sheet1']\n",
    "    for i, row in enumerate(sheet.iter_rows()):\n",
    "        if i == 0: continue\n",
    "        if row[3].value is not None and row[4].value is not None:\n",
    "            plct = row[3].value.lower().split('-')\n",
    "            pl = row[4].value.lower().split('-')\n",
    "            assert len(pl) == len(plct), f'CT和CTA标签不等长{input_path}:{patient}, CT:{plct}, CTA:{pl}'\n",
    "            \n",
    "check_label('/nfs3-p1/zsxm/dataset/2021-9-8/')\n",
    "check_label('/nfs3-p1/zsxm/dataset/2021-9-13/')\n",
    "check_label('/nfs3-p1/zsxm/dataset/2021-9-19/')\n",
    "check_label('/nfs3-p1/zsxm/dataset/2021-9-28/')\n",
    "check_label('/nfs3-p2/zsxm/dataset/2021-10-19-imh/')\n",
    "check_label('/nfs3-p2/zsxm/dataset/2021-11-20/')\n",
    "check_label('/nfs3-p2/zsxm/dataset/2021-11-20-imh/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64115f1c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 将2下的dcm文件根据窗宽窗位转化为png图片\n",
    "def generate_image(input_folder):\n",
    "    for patient in sorted(os.listdir(input_folder)):\n",
    "        if os.path.isfile(os.path.join(input_folder, patient)):\n",
    "            continue\n",
    "        print(f'****Processing {patient}****')\n",
    "        for scan in os.listdir(os.path.join(input_folder, patient)):\n",
    "            if scan != '2':\n",
    "                continue\n",
    "            name = patient #name = patient.split('-')[0]\n",
    "            image_path = os.path.join(input_folder, patient, scan, f'images_{lower_b}_{upper_b}')\n",
    "            if os.path.exists(image_path):\n",
    "                shutil.rmtree(image_path)\n",
    "            os.mkdir(image_path)\n",
    "\n",
    "            ct = load_scan(os.path.join(input_folder, patient, scan))\n",
    "            print_flag = False\n",
    "            for i in range(len(ct)):\n",
    "                img = ct[i].pixel_array.astype(np.int16)\n",
    "                intercept = ct[i].RescaleIntercept\n",
    "                slope = ct[i].RescaleSlope\n",
    "                if slope != 1:\n",
    "                    img = (slope * img.astype(np.float64)).astype(np.int16)\n",
    "                img += np.int16(intercept)\n",
    "                img = np.clip(img, lower_b, upper_b)\n",
    "                img = ((img-lower_b)/(upper_b-lower_b)*255).astype(np.uint8)\n",
    "                img = Image.fromarray(img)\n",
    "                if img.height != img.width:\n",
    "                    if not print_flag:\n",
    "                        print(patient, f'height({img.height}) not equal to width({img.width})\\n')\n",
    "                        print_flag = True\n",
    "                    height = width = min(img.height, img.width)\n",
    "                    if img.height != height:\n",
    "                        start = (img.height - height) / 2\n",
    "                        img = img.crop((0, start, img.width, start + height))\n",
    "                    elif img.width != width:\n",
    "                        start = (img.width - width) / 2\n",
    "                        img = img.crop((start, 0, start + height, img.height))\n",
    "                img.save(os.path.join(image_path, f'{name}_{i:04d}.png'))\n",
    "\n",
    "generate_image('/nfs3-p1/zsxm/dataset/2021-9-8/')\n",
    "print('----------------------------------------------------------------------------')\n",
    "generate_image('/nfs3-p1/zsxm/dataset/2021-9-13/')\n",
    "print('----------------------------------------------------------------------------')\n",
    "generate_image('/nfs3-p1/zsxm/dataset/2021-9-19/')\n",
    "print('----------------------------------------------------------------------------')\n",
    "generate_image('/nfs3-p1/zsxm/dataset/2021-9-28/')\n",
    "print('----------------------------------------------------------------------------')\n",
    "generate_image('/nfs3-p2/zsxm/dataset/2021-10-19-aa/')\n",
    "print('----------------------------------------------------------------------------')\n",
    "generate_image('/nfs3-p2/zsxm/dataset/2021-10-19-imh/')\n",
    "print('----------------------------------------------------------------------------')\n",
    "generate_image('/nfs3-p2/zsxm/dataset/2021-10-19-pau/')\n",
    "print('----------------------------------------------------------------------------')\n",
    "generate_image('/nfs3-p2/zsxm/dataset/2021-11-20/')\n",
    "print('----------------------------------------------------------------------------')\n",
    "generate_image('/nfs3-p2/zsxm/dataset/2021-11-20-imh/')\n",
    "print('----------------------------------------------------------------------------')\n",
    "generate_image('/nfs3-p2/zsxm/dataset/2021-11-20-pau/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2eee6b50",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 将各个病例中的png图片文件夹统一移动到一起供yolov5检测, not_move=True表示若有labels则不移动去检测\n",
    "def move_together_for_detect(input_folder, dst_path, not_move=True):   \n",
    "    if not os.path.exists(dst_path):\n",
    "        os.mkdir(dst_path)\n",
    "    root_name = input_folder.split('/')[-1] if input_folder.split('/')[-1] != '' else input_folder.split('/')[-2]\n",
    "    dst_path = os.path.join(dst_path, root_name)\n",
    "\n",
    "    for patient in sorted(os.listdir(input_folder)):\n",
    "        if os.path.isfile(os.path.join(input_folder, patient)):\n",
    "            continue\n",
    "        if not_move and os.path.exists(os.path.join(input_folder, patient, '2', 'labels')) \\\n",
    "        and os.path.exists(os.path.join(input_folder, patient, '2', f'pred_images_{lower_b}_{upper_b}')):\n",
    "            continue\n",
    "        print(f'****Processing {patient}****')\n",
    "        name = patient #name = patient.split('-')[0]\n",
    "        if os.path.exists(os.path.join(dst_path, name)):\n",
    "            print(f\"\\tremove {os.path.join(dst_path, name)}\")\n",
    "            shutil.rmtree(os.path.join(dst_path, name))\n",
    "\n",
    "        try:\n",
    "            shutil.copytree(os.path.join(input_folder, patient, '2', f'images_{lower_b}_{upper_b}'), os.path.join(dst_path, name))\n",
    "        except:\n",
    "            traceback.print_exc()\n",
    "\n",
    "move_together_for_detect('/nfs3-p1/zsxm/dataset/2021-9-8/', '/nfs3-p1/zsxm/dataset/9_detect/')\n",
    "move_together_for_detect('/nfs3-p1/zsxm/dataset/2021-9-13/', '/nfs3-p1/zsxm/dataset/9_detect/')\n",
    "move_together_for_detect('/nfs3-p1/zsxm/dataset/2021-9-19/', '/nfs3-p1/zsxm/dataset/9_detect/')\n",
    "move_together_for_detect('/nfs3-p1/zsxm/dataset/2021-9-28/', '/nfs3-p1/zsxm/dataset/9_detect/')\n",
    "#move_together_for_detect('/nfs3-p2/zsxm/dataset/2021-10-19-aa/', '/nfs3-p1/zsxm/dataset/9_detect/')\n",
    "move_together_for_detect('/nfs3-p2/zsxm/dataset/2021-10-19-imh/', '/nfs3-p1/zsxm/dataset/9_detect/')\n",
    "#move_together_for_detect('/nfs3-p2/zsxm/dataset/2021-10-19-pau/', '/nfs3-p1/zsxm/dataset/9_detect/')\n",
    "move_together_for_detect('/nfs3-p2/zsxm/dataset/2021-11-20/', '/nfs3-p1/zsxm/dataset/9_detect/')\n",
    "move_together_for_detect('/nfs3-p2/zsxm/dataset/2021-11-20-imh/', '/nfs3-p1/zsxm/dataset/9_detect/')\n",
    "#move_together_for_detect('/nfs3-p2/zsxm/dataset/2021-11-20-pau/', '/nfs3-p1/zsxm/dataset/9_detect/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1761401d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#将检测结果移动回原文件夹内\n",
    "def move_back(result_path, ori_path):\n",
    "    if not os.path.exists(result_path):\n",
    "        print(f'目录不存在：{result_path}')\n",
    "        return\n",
    "    for patient in sorted(os.listdir(result_path)):\n",
    "        print(f'Processing {patient}')\n",
    "        p_res_path = os.path.join(result_path, patient)\n",
    "        o_res_path = os.path.join(ori_path, patient, '2', f'pred_images_{lower_b}_{upper_b}')\n",
    "        if os.path.exists(o_res_path):\n",
    "            shutil.rmtree(o_res_path)\n",
    "        os.mkdir(o_res_path)\n",
    "        for file in os.listdir(p_res_path):\n",
    "            if os.path.isfile(os.path.join(p_res_path, file)):\n",
    "                shutil.move(os.path.join(p_res_path, file), os.path.join(o_res_path, file))\n",
    "            elif os.path.isdir(os.path.join(p_res_path, file)):\n",
    "                if os.path.exists(os.path.join(ori_path, patient, '2', file)):\n",
    "                    shutil.rmtree(os.path.join(ori_path, patient, '2', file))\n",
    "                shutil.move(os.path.join(p_res_path, file), os.path.join(ori_path, patient, '2', file))\n",
    "        os.rmdir(p_res_path)\n",
    "    os.rmdir(result_path)\n",
    "\n",
    "\n",
    "#move_back('/home/zsxm/pythonWorkspace/yolov5_old/runs/detect/2021-9-8', '/nfs3-p2/zsxm/dataset/2021-9-8/')\n",
    "move_back('/home/zsxm/pythonWorkspace/yolov5_old/runs/detect/2021-9-13', '/nfs3-p1/zsxm/dataset/2021-9-13/')\n",
    "move_back('/home/zsxm/pythonWorkspace/yolov5_old/runs/detect/2021-9-19', '/nfs3-p2/zsxm/dataset/2021-9-19/')\n",
    "move_back('/home/zsxm/pythonWorkspace/yolov5_old/runs/detect/2021-9-28', '/nfs3-p2/zsxm/dataset/2021-9-28/')\n",
    "#move_back('/home/zsxm/pythonWorkspace/yolov5_old/runs/detect/2021-10-19-imh', '/nfs3-p2/zsxm/dataset/2021-10-19-imh/')\n",
    "move_back('/home/zsxm/pythonWorkspace/yolov5_old/runs/detect/2021-11-20', '/nfs3-p1/zsxm/dataset/2021-11-20/')\n",
    "move_back('/home/zsxm/pythonWorkspace/yolov5_old/runs/detect/2021-11-20-imh', '/nfs3-p2/zsxm/dataset/2021-11-20-imh/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "574510e4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 切出主动脉\n",
    "def find_coordinate(height, width, label_file, aorta):\n",
    "    with open(label_file, 'r') as f:\n",
    "        lines = f.readlines()\n",
    "    assert len(lines) <= 2, f'label.txt应该存储不多于2个label：{label_file.split(\"/\")[-1]}'\n",
    "    if len(lines) == 1:\n",
    "        assert aorta == 'j', f'如果只有一个label那么此时应为降主动脉, 但实际为{aorta}：{label_file.split(\"/\")[-1]}'\n",
    "        corr = list(map(lambda x: float(x), lines[0].split()))\n",
    "        x, y, w, h = corr[1], corr[2], corr[3], corr[4]\n",
    "        assert 0.25 < x < 0.75 and 0.15 < y < 0.85, f'边界框中心({x}, {y})出界：{label_file.split(\"/\")[-1]}'\n",
    "    else:\n",
    "        corr1, corr2 = list(map(lambda x: float(x), lines[0].split())), list(map(lambda x: float(x), lines[1].split()))\n",
    "        assert 0.25 < corr1[1] < 0.75 and 0.15 < corr1[2] < 0.85, f'边界框1中心({corr1[1]}, {corr1[2]})出界：{label_file.split(\"/\")[-1]}'\n",
    "        assert 0.25 < corr2[1] < 0.75 and 0.15 < corr2[2] < 0.85, f'边界框2中心({corr2[1]}, {corr2[2]})出界：{label_file.split(\"/\")[-1]}'\n",
    "        if aorta == 's':\n",
    "            x, y, w, h = (corr1[1], corr1[2], corr1[3], corr1[4]) if corr1[2] < corr2[2] else (corr2[1], corr2[2], corr2[3], corr2[4])\n",
    "        elif aorta == 'j':\n",
    "            x, y, w, h = (corr1[1], corr1[2], corr1[3], corr1[4]) if corr1[2] > corr2[2] else (corr2[1], corr2[2], corr2[3], corr2[4])\n",
    "        else:\n",
    "            raise Exception(f'aorta 应该为\"s\"或\"j\"其中之一: {label_file.split(\"/\")[-1]}')\n",
    "    w, h = int(width*w), int(height*h)\n",
    "    w, h = max(w, h), max(w, h)\n",
    "    return int(width*x-w/2), int(height*y-h/2), int(width*x+w/2+1), int(height*y+h/2+1)\n",
    "\n",
    "def crop_images(input_path, error_patient_list):\n",
    "    workbook_path = os.path.join(input_path, 'label.xlsx')\n",
    "    wb = openpyxl.load_workbook(workbook_path)\n",
    "    sheet = wb['Sheet1']\n",
    "    \n",
    "    for patient in sorted(os.listdir(input_path)):\n",
    "        if os.path.isfile(os.path.join(input_path, patient)):\n",
    "            continue\n",
    "        flag = True\n",
    "        for row in sheet.iter_rows():\n",
    "            if row[0].value == patient.split('-')[0]:\n",
    "                if row[3].value is not None and row[4].value is not None:\n",
    "                    flag = False\n",
    "                    pl = row[4].value.lower().split('-')\n",
    "                    plct = row[3].value.lower().split('-')\n",
    "                    assert len(pl) == len(plct), f'CT和CTA标签不等长{input_path}:{patient}, {len(pl)}, {len(plct)}'\n",
    "                break\n",
    "        if flag: continue\n",
    "        print(f'******Processing {patient}******')\n",
    "        image_path = os.path.join(input_path, patient, '2', f'images_{lower_b}_{upper_b}')\n",
    "        label_path = os.path.join(input_path, patient, '2', 'labels')\n",
    "        crop_path = os.path.join(input_path, patient, '2', f'crops_{lower_b}_{upper_b}')\n",
    "        if os.path.exists(crop_path):\n",
    "            shutil.rmtree(crop_path)\n",
    "        os.mkdir(crop_path)\n",
    "        \n",
    "        crop_flag = True\n",
    "        for i, s in enumerate(pl):\n",
    "            if s != 's' and s != 'j':\n",
    "                continue\n",
    "            start, end = int(pl[i+1])-1, int(pl[i+2])\n",
    "            for j in range(start, end):\n",
    "                img = Image.open(os.path.join(image_path, f'{patient}_{j:04d}.png'))\n",
    "                img = np.array(img)\n",
    "                try:\n",
    "                    x1, y1, x2, y2 = find_coordinate(*img.shape[0:2], os.path.join(label_path, f'{patient}_{j:04d}.txt'), s)\n",
    "                except:\n",
    "                    traceback.print_exc()\n",
    "                    crop_flag = False\n",
    "                else:#if crop_flag:\n",
    "                    crop = img[y1:y2, x1:x2]\n",
    "                    crop = Image.fromarray(crop)\n",
    "                    crop.save(os.path.join(crop_path, f'{patient}_{s}_{j:04d}.png'))\n",
    "        if not crop_flag:\n",
    "            error_patient_list.append(patient)\n",
    "\n",
    "epl1 = []\n",
    "crop_images('/nfs3-p1/zsxm/dataset/2021-9-8/', epl1)\n",
    "crop_images('/nfs3-p1/zsxm/dataset/2021-9-13/', epl1)\n",
    "crop_images('/nfs3-p1/zsxm/dataset/2021-9-19/', epl1)\n",
    "crop_images('/nfs3-p1/zsxm/dataset/2021-9-28/', epl1)\n",
    "crop_images('/nfs3-p1/zsxm/dataset/2021-10-19-imh/', epl1)\n",
    "crop_images('/nfs3-p1/zsxm/dataset/2021-11-20/', epl1)\n",
    "crop_images('/nfs3-p1/zsxm/dataset/2021-11-20-imh/', epl1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0df154bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(epl1))\n",
    "print(epl1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5283cfa9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# 切出范围外冗余为3的主动脉\n",
    "def find_coordinate(height, width, label_file, aorta):\n",
    "    with open(label_file, 'r') as f:\n",
    "        lines = f.readlines()\n",
    "    assert len(lines) <= 2, f'label.txt应该存储不多于2个label：{label_file.split(\"/\")[-1]}'\n",
    "    if len(lines) == 1:\n",
    "        assert aorta == 'j', f'如果只有一个label那么此时应为降主动脉, 但实际为{aorta}：{label_file.split(\"/\")[-1]}'\n",
    "        corr = list(map(lambda x: float(x), lines[0].split()))\n",
    "        x, y, w, h = corr[1], corr[2], corr[3], corr[4]\n",
    "        assert 0.25 < x < 0.75 and 0.15 < y < 0.85, f'边界框中心({x}, {y})出界：{label_file.split(\"/\")[-1]}'\n",
    "    else:\n",
    "        corr1, corr2 = list(map(lambda x: float(x), lines[0].split())), list(map(lambda x: float(x), lines[1].split()))\n",
    "        assert 0.25 < corr1[1] < 0.75 and 0.15 < corr1[2] < 0.85, f'边界框1中心({corr1[1]}, {corr1[2]})出界：{label_file.split(\"/\")[-1]}'\n",
    "        assert 0.25 < corr2[1] < 0.75 and 0.15 < corr2[2] < 0.85, f'边界框2中心({corr2[1]}, {corr2[2]})出界：{label_file.split(\"/\")[-1]}'\n",
    "        if aorta == 's':\n",
    "            x, y, w, h = (corr1[1], corr1[2], corr1[3], corr1[4]) if corr1[2] < corr2[2] else (corr2[1], corr2[2], corr2[3], corr2[4])\n",
    "        elif aorta == 'j':\n",
    "            x, y, w, h = (corr1[1], corr1[2], corr1[3], corr1[4]) if corr1[2] > corr2[2] else (corr2[1], corr2[2], corr2[3], corr2[4])\n",
    "        else:\n",
    "            raise Exception(f'aorta 应该为\"s\"或\"j\"其中之一: {label_file.split(\"/\")[-1]}')\n",
    "    w, h = int(width*w), int(height*h)\n",
    "    w, h = max(w, h), max(w, h)\n",
    "    return int(width*x-w/2), int(height*y-h/2), int(width*x+w/2+1), int(height*y+h/2+1)\n",
    "\n",
    "def crop_images(input_path, error_patient_list):\n",
    "    workbook_path = os.path.join(input_path, 'label.xlsx')\n",
    "    wb = openpyxl.load_workbook(workbook_path)\n",
    "    sheet = wb['Sheet1']\n",
    "    \n",
    "    for patient in sorted(os.listdir(input_path)):\n",
    "        if os.path.isfile(os.path.join(input_path, patient)):\n",
    "            continue\n",
    "        flag = True\n",
    "        for row in sheet.iter_rows():\n",
    "            if row[0].value == patient.split('-')[0]:\n",
    "                if row[3].value is not None and row[4].value is not None:\n",
    "                    flag = False\n",
    "                    pl = row[4].value.lower().split('-')\n",
    "                    plct = row[3].value.lower().split('-')\n",
    "                    assert len(pl) == len(plct), f'CT和CTA标签不等长{input_path}:{patient}, {len(pl)}, {len(plct)}'\n",
    "                break\n",
    "        if flag: continue\n",
    "        \n",
    "        print(f'******Processing {patient}******')\n",
    "        image_path = os.path.join(input_path, patient, '2', f'images_{lower_b}_{upper_b}')\n",
    "        label_path = os.path.join(input_path, patient, '2', 'labels')\n",
    "        crop_path = os.path.join(input_path, patient, '2', f'crops3_{lower_b}_{upper_b}')\n",
    "        if os.path.exists(crop_path):\n",
    "            shutil.rmtree(crop_path)\n",
    "        os.mkdir(crop_path)\n",
    "        \n",
    "        crop_flag = True\n",
    "        for i, s in enumerate(pl):\n",
    "            if s != 's' and s != 'j':\n",
    "                continue\n",
    "            start, end = int(pl[i+1])-1, int(pl[i+2])\n",
    "            for j in range(start-3, end+3):\n",
    "                try:\n",
    "                    img = Image.open(os.path.join(image_path, f'{patient}_{j:04d}.png'))\n",
    "                    img = np.array(img)\n",
    "                    x1, y1, x2, y2 = find_coordinate(*img.shape[0:2], os.path.join(label_path, f'{patient}_{j:04d}.txt'), s)\n",
    "                except:\n",
    "                    traceback.print_exc()\n",
    "                    crop_flag = False\n",
    "                else:#if crop_flag:\n",
    "                    crop = img[y1:y2, x1:x2]\n",
    "                    crop = Image.fromarray(crop)\n",
    "                    if start <= j < end:\n",
    "                        crop.save(os.path.join(crop_path, f'{patient}_{s}_{j:04d}.png'))\n",
    "                    else:\n",
    "                        crop.save(os.path.join(crop_path, f'{patient}_{s}_{j:04d}_n.png'))\n",
    "        if not crop_flag:\n",
    "            error_patient_list.append(patient)\n",
    "            \n",
    "epl1 = []\n",
    "crop_images('/nfs3-p1/zsxm/dataset/2021-9-8/', epl1)\n",
    "crop_images('/nfs3-p1/zsxm/dataset/2021-9-13/', epl1)\n",
    "crop_images('/nfs3-p1/zsxm/dataset/2021-9-19/', epl1)\n",
    "crop_images('/nfs3-p1/zsxm/dataset/2021-9-28/', epl1)\n",
    "crop_images('/nfs3-p1/zsxm/dataset/2021-10-19-imh/', epl1)\n",
    "crop_images('/nfs3-p1/zsxm/dataset/2021-11-20/', epl1)\n",
    "crop_images('/nfs3-p1/zsxm/dataset/2021-11-20-imh/', epl1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5419ed51",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(epl1))\n",
    "print(epl1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a08fa9d",
   "metadata": {},
   "source": [
    "# 3.复制文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "943572f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "classify_path = f'/nfs3-p2/zsxm/dataset/aorta_classify_cta_{lower_b}_{upper_b}'\n",
    "os.makedirs(classify_path, exist_ok=True)\n",
    "for dataset in ['train', 'val']:\n",
    "    dst_path = os.path.join(classify_path, dataset)\n",
    "    os.makedirs(dst_path, exist_ok=True)\n",
    "    for cate in range(3):\n",
    "        cls_path = os.path.join(dst_path, str(cate))\n",
    "        os.makedirs(cls_path, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "370d1c04",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set = set()\n",
    "val_set = set()\n",
    "ct_path = f'/nfs3-p2/zsxm/dataset/aorta_classify_ct_{lower_b}_{upper_b}/'\n",
    "for cate in os.listdir(os.path.join(ct_path, 'train')):\n",
    "    for img in os.listdir(os.path.join(ct_path, 'train', cate)):\n",
    "        train_set.add(img.split('_')[0])\n",
    "for img in os.listdir(os.path.join(ct_path, 'val')):\n",
    "    for img in os.listdir(os.path.join(ct_path, 'val', cate)):\n",
    "        val_set.add(img.split('_')[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acb97424",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(train_set), len(val_set))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c28706a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def move_cta(input_path, cate, train_set, val_set):\n",
    "    workbook_path = os.path.join(input_path, 'label.xlsx')\n",
    "    wb = openpyxl.load_workbook(workbook_path)\n",
    "    sheet = wb['Sheet1']\n",
    "    \n",
    "    for patient in sorted(os.listdir(input_path)):\n",
    "        if os.path.isfile(os.path.join(input_path, patient)):\n",
    "            continue\n",
    "        flag = True\n",
    "        for row in sheet.iter_rows():\n",
    "            if row[0].value == patient.split('-')[0]:\n",
    "                if row[3].value is not None and row[4].value is not None:\n",
    "                    flag = False\n",
    "                break\n",
    "        if flag: continue\n",
    "        print(f'******Processing {patient}******')\n",
    "        if patient in train_set:\n",
    "            dst_path = os.path.join(classify_path, 'train', str(cate))\n",
    "        elif patient in val_set:\n",
    "            dst_path = os.path.join(classify_path, 'val', str(cate))\n",
    "        else:\n",
    "            raise Exception(f'{patient} neither in train_set nor in val_set')\n",
    "        ori_path = os.path.join(input_path, patient, '2', f'crops_{lower_b}_{upper_b}')\n",
    "        for img in os.listdir(ori_path):\n",
    "            shutil.copy(os.path.join(ori_path, img), os.path.join(dst_path, img))\n",
    "            \n",
    "\n",
    "move_cta('/nfs3-p1/zsxm/dataset/2021-9-17-negative/', train_set, val_set, 0)\n",
    "move_cta('/nfs3-p1/zsxm/dataset/2021-9-29-negative/', train_set, val_set, 0)\n",
    "move_cta('/nfs3-p1/zsxm/dataset/2021-9-8/', train_set, val_set, 1)\n",
    "move_cta('/nfs3-p1/zsxm/dataset/2021-9-13/', train_set, val_set, 1)\n",
    "move_cta('/nfs3-p1/zsxm/dataset/2021-9-19/', train_set, val_set, 1)\n",
    "move_cta('/nfs3-p1/zsxm/dataset/2021-9-28/', train_set, val_set, 1)\n",
    "move_cta('/nfs3-p1/zsxm/dataset/2021-10-19-imh/', train_set, val_set, 2)\n",
    "move_cta('/nfs3-p1/zsxm/dataset/2021-11-20/', train_set, val_set, 1)\n",
    "move_cta('/nfs3-p1/zsxm/dataset/2021-11-20-imh/', train_set, val_set, 2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
